from dataclasses import dataclass
from typing import Tuple, Sequence, Any, Union, Optional
import typing as tp

import numpy as np
import pandas as pd

from . import base
from ..multiprocessing import mpPandasObj


@dataclass
class TripleGrid(base.BaseTransformer, base.CopyMixin):
    """栅格法（三分类）"""
    n: int = 5
    upper: float = 0.005
    lower: float = -0.005
    new_col: str = 'y'
    col: str = 'close'
    labels: Tuple[int, int, int] = (1, 0, -1)  # 涨/平/跌

    @classmethod
    def get_triple_grid(cls,
                        x: pd.Series,
                        upper: float,
                        lower: float,
                        labels: Tuple[int, int, int] = (1, 0, -1)):
        up, end, low = labels
        base_ = x.iloc[0]
        for i in x:
            if i / base_ - 1 <= lower:
                return low
            if i / base_ - 1 >= upper:
                return up
        return end

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = data[self.col].rolling(self.n + 1).agg(
            lambda x: self.get_triple_grid(x, self.upper, self.lower, labels=self.labels)
        ).shift(-self.n - 1)
        return data


@dataclass
class QuadraGrid(base.BaseTransformer, base.CopyMixin):
    """栅格法（四分类）"""
    n: int = 5
    upper: float = 0.005
    lower: float = -0.005
    new_col: str = 'y'
    col: str = 'close'
    labels: Tuple[int, int, int, int] = (2, 1, -1, -2)  # 涨/微涨/微跌/跌

    @classmethod
    def get_quadra_grid(cls,
                        x: pd.Series,
                        upper: float,
                        lower: float,
                        labels: Tuple[int, int, int, int] = (2, 1, -1, -2)):
        up, uu, ll, low = labels
        base_ = x.iloc[0]
        for i in x:
            if i / base_ - 1 <= lower:
                return low
            if i / base_ - 1 >= upper:
                return up
        return uu if x.iloc[-1] / base_ - 1 > 0 else ll

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = data[self.col].rolling(self.n + 1).agg(
            lambda x: self.get_quadra_grid(x, self.upper, self.lower, labels=self.labels)
        ).shift(-self.n - 1)
        return data


@dataclass
class DynTripleGrid(base.BaseTransformer, base.CopyMixin):
    n: int = 5
    span: int = 20
    alpha_u: float = 1.
    alpha_d: float = 1.
    new_col: str = 'y'
    col: str = 'close'
    labels: Tuple[int, int, int] = (1, 0, -1)  # 涨/平/跌

    @classmethod
    def get_triple_grid(cls,
                        x: pd.Series,
                        y: pd.Series,
                        alpha_u: float,
                        alpha_d: float,
                        labels: Tuple[int, int, int] = (1, 0, -1)):
        up, end, low = labels
        base_ = x.iloc[0]
        std = y.loc[x.index].iloc[0]
        for i in x.iloc[1:]:
            lower = - alpha_d * std
            upper = alpha_u * std
            if i / base_ - 1 <= lower:
                return low
            if i / base_ - 1 >= upper:
                return up
        return end

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        ret = data[self.col].pct_change()
        stds = ret.ewm(span=self.span).std()
        data[self.new_col] = data[self.col].rolling(self.n + 1).agg(
            lambda x: self.get_triple_grid(x, stds, self.alpha_u, self.alpha_d, labels=self.labels)
        ).shift(-self.n - 1)
        return data


@dataclass
class PnLRatio(base.BaseTransformer, base.CopyMixin):
    """盈亏比"""
    new_col: str = "pnl_ratio"
    col: str = 'close'
    n: int = 10
    long: bool = True
    target_ratio: float = 5.

    @classmethod
    def trans(cls, x: pd.Series, n: int, long: bool):
        roll_max = x.rolling(n).max().shift(-n)
        roll_min = x.rolling(n).min().shift(-n)
        inc = roll_max / x - 1 + 1 / 100000  # 防止等于0
        dec = roll_min / x - 1 + 1 / 100000
        ratio = inc / dec if long else dec / inc
        return ratio.abs().shift(-n)

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = self.trans(data[self.col], self.n, self.long)
        if self.target_ratio > 0:
            data[self.new_col] = (data[self.new_col] >= self.target_ratio).astype('int')
        return data


@dataclass
class MaxMinGrid(base.BaseTransformer, base.CopyMixin):
    new_col: str = "label"
    col: str = "close"
    col_upper: str = "high"
    col_lower: str = "low"
    n: int = 1
    upper: float = 0.002
    lower: float = -0.002
    labels: Tuple[int, int, int] = (1, 0, -1)  # 涨/平/跌

    @classmethod
    def get_triple_grid(cls,
                        x: pd.Series,
                        a: pd.Series,
                        b: pd.Series,
                        upper: float,
                        lower: float,
                        labels: Tuple[int, int, int] = (1, 0, -1)):
        up, end, low = labels
        base_ = x.iloc[0]
        for u, l in zip(a.loc[x.iloc[1:].index], b.loc[x.iloc[1:].index]):
            if l / base_ - 1 <= lower:
                return low
            if u / base_ - 1 >= upper:
                return up
        return end

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = data[self.col].rolling(self.n + 1).agg(
            lambda x: self.get_triple_grid(x,
                                           data[self.col_upper],
                                           data[self.col_lower],
                                           self.upper,
                                           self.lower,
                                           labels=self.labels)
        ).shift(-self.n - 1)
        return data


def get_t1(close: pd.Series, 
           t_sample: pd.Index,
           value: int = 1, 
           unit: str = "days"):
    """指定时间间隔计算数值隔栏的时间戳
    
    :param close: 收盘价, index为datetime
    :param t_sample: 抽样后的index
    :param value: 计算Timedelta的value
    :param unit: 计算Timedelta的unit
    """
    t1 = close.index.searchsorted(t_sample+pd.Timedelta(value=value, unit=unit))
    t1 = t1[t1 < close.shape[0]]
    t1 = pd.Series(close.index[t1],
                   index=t_sample[:t1.shape[0]])  # NaNs at end
    return t1


def applyPtSlOnT1(close: pd.Series,
                  events: pd.DataFrame,
                  ptSl: Tuple[float, float],
                  molecule: Sequence[Any]):
    """计算触碰三隔栏的时间戳
    
    `events`要有两列：
    * t1: 垂直隔栏时间戳
    * trgt: 水平隔栏的单位宽度

    :param close: A pandas series of prices.
    :param events: A pandas dataframe, with columns
    :param ptSl: A list of two non-negative float values:
    :param molecule:
    :return:
    """
    # apply stop loss/profit taking, if it takes place before t1 (end of event)
    events_ = events.loc[molecule]
    out = events_[['t1']].copy(deep=True)
    if ptSl[0] > 0:
        pt = ptSl[0] * events_['trgt']
    else:
        pt = pd.Series(index=events.index)  # NaNs
    if ptSl[1] > 0:
        sl = -ptSl[1] * events_['trgt']
    else:
        sl = pd.Series(index=events.index)  # NaNs
    for loc, t1 in events_['t1'].fillna(close.index[-1]).iteritems():
        df0 = close[loc:t1]  # path prices
        df0 = (df0 / close[loc] - 1) * events_.at[loc, 'side']  # path returns
        out.loc[loc, 'sl'] = df0[df0 < sl[loc]].index.min()  # earliest stop loss.
        out.loc[loc, 'pt'] = df0[df0 > pt[loc]].index.min()  # earliest profit taking.
    return out


def getEvents(close: pd.Series,
              tEvents: pd.Index,
              ptSl: Tuple[float, float],
              trgt: pd.Series,
              minRet: float,
              numThreads: int,
              t1: Optional[pd.Series] = None,
              side: Optional[pd.Series] = None):
    """计算第一次触碰隔栏的时间戳

    :param close: 收盘价
    :param tEvents: 抽样后的time-index
    :param ptSl:  上下隔栏的宽度，非负数
    :param trgt:  绝对回报
    :param minRet:  最小回报
    :param numThreads: 线程数
    :param t1: 垂直栏的时间戳
    :param side: 头寸方向
    :return:
    """
    # 1) get target
    trgt = trgt.loc[tEvents]
    trgt = trgt[trgt > minRet]  # minRet
    # 2) get t1 (max holding period)
    if t1 is None:
        t1 = pd.Series(pd.NaT, index=tEvents)
    # 3) form events object, apply stop loss on t1
    if side is None:
        side_, ptSl_ = pd.Series(1., index=trgt.index), (ptSl[0], ptSl[0])
    else:
        side_, ptSl_ = side.loc[trgt.index], ptSl[:2]
    events = pd.concat({'t1': t1, 'trgt': trgt, 'side': side_}, axis=1)
    events = tp.cast(pd.DataFrame, events)
    events = events.dropna(subset=['trgt'])
    df0 = applyPtSlOnT1(close=close, events=events, ptSl=ptSl_, molecule=events.index)
    # 获取第一次触碰隔栏的时间
    events['t1'] = df0.dropna(how='all').min(axis=1)  # pd.min ignores nan
    if side is None:
        events = events.drop('side', axis=1)
    return events


def getEventsMulti(close: pd.Series,
              tEvents: pd.Index,
              ptSl: Tuple[float, float],
              trgt: pd.Series,
              minRet: float,
              numThreads: int,
              t1: Optional[pd.Series] = None,
              side: Optional[pd.Series] = None):
    """计算第一次触碰隔栏的时间戳

    :param close: 收盘价
    :param tEvents: 抽样后的time-index
    :param ptSl:  上下隔栏的宽度，非负数
    :param trgt:  绝对回报
    :param minRet:  最小回报
    :param numThreads: 线程数
    :param t1: 垂直栏的时间戳
    :param side: 头寸方向
    :return:
    """
    # 1) get target
    trgt = trgt.loc[tEvents]
    trgt = trgt[trgt > minRet]  # minRet
    # 2) get t1 (max holding period)
    if t1 is None:
        t1 = pd.Series(pd.NaT, index=tEvents)
    # 3) form events object, apply stop loss on t1
    if side is None:
        side_, ptSl_ = pd.Series(1., index=trgt.index), [ptSl[0], ptSl[0]]
    else:
        side_, ptSl_ = side.loc[trgt.index], ptSl[:2]
    events = pd.concat({'t1': t1, 'trgt': trgt, 'side': side_}, axis=1)
    events = tp.cast(pd.DataFrame, events)
    events = events.dropna(subset=['trgt'])
    df0 = mpPandasObj(func=applyPtSlOnT1, pdObj=('molecule', events.index),
                      numThreads=numThreads, close=close, events=events, ptSl=ptSl_)
    df0 = tp.cast(pd.DataFrame, df0)
    # 获取第一次触碰隔栏的时间
    events['t1'] = df0.dropna(how='all').min(axis=1)  # pd.min ignores nan
    if side is None:
        events = events.drop('side', axis=1)
    return events


def getBins(events: pd.DataFrame, close: pd.Series, is_bin=False):
    """为下注方向和大小打标签

    Compute event's outcome (including side information, if provided).
    events is a DataFrame where:
    —events.index is event's starttime
    —events[’t1’] is event's endtime
    —events[’trgt’] is event's target
    —events[’side’] (optional) implies the algo's position side，`side`是主模型预测的方向
    Case 1: (’side’ not in events): bin in (-1,1) <—label by price action
    Case 2: (’side’ in events): bin in (0,1) <—label by pnl (meta-labeling)


    :param events: 事件
    :param close: 收盘价
    :return:
    """
    # 1) prices aligned with events
    events_ = events.dropna(subset=['t1'])
    px = events_.index.union(events_['t1'].values).drop_duplicates()
    px = close.reindex(px, method='bfill')
    # 2) create out object
    out = pd.DataFrame(index=events_.index)
    out['ret'] = px.loc[events_['t1'].values].values / px.loc[events_.index] - 1
    # 如果events中包含主模型预测的方向
    # 就修改ret为主模型预测方向的收益率
    if "side" in events_:
        out['ret'] *= events_['side']  # meta-labeling
    # out['bin'] = np.sign(out['ret'])
    # 打标签
    if is_bin:
        out["bin"] = np.where(out['ret'] < 0, -1, 1)
    else:
        out['bin'] = np.select([out['ret'] < 0, 
                                out['ret'] == 0, 
                                out['ret'] > 0],
                               [-1, 0, 1])
    if "side" in events_:
        # 主模型预测错的方向标签设置为0
        out.loc[out['ret'] <= 0, 'bin'] = 0  # meta-labeling
    return out
